package com.cy.deepseeksport;

import io.github.pigmesh.ai.deepseek.core.shared.Usage;

import java.time.LocalTime;

public class DeepSeekCostCalculator {

    // 标准时段价格（元/百万tokens）
    private static final double STANDARD_CACHE_HIT_INPUT = 0.5;
    private static final double STANDARD_CACHE_MISS_INPUT = 2.0;
    private static final double STANDARD_OUTPUT = 8.0;

    // 优惠时段价格（元/百万tokens）
    private static final double DISCOUNT_CACHE_HIT_INPUT = 0.25;
    private static final double DISCOUNT_CACHE_MISS_INPUT = 1.0;
    private static final double DISCOUNT_OUTPUT = 4.0;

    // 时段定义
    private static final LocalTime DISCOUNT_START = LocalTime.of(0, 30);
    private static final LocalTime DISCOUNT_END = LocalTime.of(8, 30);


    public static double calculateCost(Usage usage) {
        if (usage == null) {
            throw new IllegalArgumentException("Usage cannot be null");
        }

        boolean isDiscountPeriod = isDiscountPeriod();
        boolean isCacheHit = isCacheHit(usage);

        // 计算输入token成本（考虑缓存命中情况）
        double inputCost;
        if (isDiscountPeriod) {
            inputCost = isCacheHit ?
                    (usage.promptTokens() / 1_000_000.0) * DISCOUNT_CACHE_HIT_INPUT :
                    (usage.promptTokens() / 1_000_000.0) * DISCOUNT_CACHE_MISS_INPUT;
        } else {
            inputCost = isCacheHit ?
                    (usage.promptTokens() / 1_000_000.0) * STANDARD_CACHE_HIT_INPUT :
                    (usage.promptTokens() / 1_000_000.0) * STANDARD_CACHE_MISS_INPUT;
        }

        // 计算输出token成本
        double outputCost;
        if (isDiscountPeriod) {
            outputCost = (usage.completionTokens() / 1_000_000.0) * DISCOUNT_OUTPUT;
        } else {
            outputCost = (usage.completionTokens() / 1_000_000.0) * STANDARD_OUTPUT;
        }

        return inputCost + outputCost;
    }

    private static boolean isCacheHit(Usage usage) {
        if (usage.promptTokensDetails() == null) {
            return false;
        }
        // 假设PromptTokensDetails中有cachedTokens字段
        // 如果有缓存的tokens，则认为有缓存命中
        return usage.promptTokensDetails().cachedTokens() != null
                && usage.promptTokensDetails().cachedTokens() > 0;
    }

    private static boolean isDiscountPeriod() {
        LocalTime now = LocalTime.now();
        // 优惠时段是00:30-08:30
        if (DISCOUNT_START.isBefore(DISCOUNT_END)) {
            return !now.isBefore(DISCOUNT_START) && now.isBefore(DISCOUNT_END);
        } else {
            return !now.isBefore(DISCOUNT_START) || now.isBefore(DISCOUNT_END);
        }
    }
}


